import taichi as ti
import random
import numpy as np

@ti.data_oriented
class JaccobiSolver:
    
    def __init__(self,A: ti.template(), b: ti.template(), epochs: ti.i32):
        shape = A.shape
        self.n = shape[1]
        self.epochs = epochs
        self.A = A
        self.x = ti.field(dtype=ti.f32, shape=shape[1])
        self.new_x = ti.field(dtype=ti.f32, shape=shape[1])
        self.b = b

    @ti.kernel
    def iterate(self):
        for i in range(self.n):
            r = self.b[i]
            for j in range(self.n):
                if i != j:
                    r -= self.A[i, j] * self.x[j]   
            self.new_x[i] = r / self.A[i, i]
        for i in range(self.n):
            self.x[i] = self.new_x[i]

    ## 计算误差
    @ti.kernel
    def residual(self) -> ti.f32:
        res = 0.0
        # 误差向量e=b-A*x
        for i in range(self.n):
            r = self.b[i] * 1.0
            for j in range(self.n):
                r -= self.A[i, j] * self.x[j]
            res += r * r
        # res实际上是e的长度的平方    
        return res
    
    def solve(self):
        ## 100次迭代
        for i in range(self.epochs):
            self.iterate()
            print(f'iter {i}, residual={self.residual():0.10f}')
        return self.x


if __name__ == '__main__':
    help(ti.SparseSolver)
    ti.init(arch = ti.gpu)
    A_np = np.array([[4,2,1],[2,5,-2],[1,-2,7]])
    b_np = np.array([23,29,0])
    A = ti.field(ti.f32,shape=(3,3))
    A.from_numpy(A_np)
    b = ti.field(ti.f32,shape=3)
    b.from_numpy(b_np)

    solver = JaccobiSolver(A=A,b=b,epochs=35)
    res = solver.solve()
    print(res)